# main.py
import sys
import os
import math
from datetime import datetime
import pandas as pd
from openpyxl import load_workbook
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QPushButton, QFileDialog, QMessageBox, QProgressBar, 
                            QTextEdit, QSplitter, QLineEdit, QComboBox, QFormLayout)
from PyQt5.QtCore import Qt, QThread, pyqtSignal


class BomProcessor(QThread):
    progress_updated = pyqtSignal(int)
    message_updated = pyqtSignal(str)
    finished_signal = pyqtSignal(str, bool)
    error_occurred = pyqtSignal(str)

    def __init__(self, source_path, template_path, output_dir, top_nodes="", split_size=0):
        super().__init__()
        self.source_path = source_path
        self.template_path = template_path
        self.output_dir = output_dir
        self.top_nodes = [x.strip() for x in top_nodes.split(",") if x.strip()] if top_nodes else []
        self.split_size = split_size
        self.max_excel_rows = 1048576  # Excel最大行数限制

    def process_dates(self, df):
        """处理日期列，固定终止时间为2999-12-31"""
        if "有效起始时间" in df.columns:
            try:
                # 处理起始时间
                df["有效起始时间"] = pd.to_datetime(df["有效起始时间"], errors='coerce')
                df["有效起始时间"] = df["有效起始时间"].apply(
                    lambda x: x.strftime("%Y-%m-%d") if not pd.isnull(x) else ""
                )
            except pd.errors.OutOfBoundsDatetime as e:
                error_msg = f"起始时间转换错误: {str(e)}"
                print(error_msg)
                self.error_occurred.emit(error_msg)
                df["有效起始时间"] = df["有效起始时间"].astype(str)
        
        # 固定终止时间为2999-12-31
        df["有效终止时间"] = "2999-12-31"
        return df

    def save_to_excel(self, data, file_path):
        """将数据保存到Excel模板文件"""
        try:
            # 加载模板文件
            book = load_workbook(self.template_path)
            sheet = book.active

            # 清除模板中现有数据（保留表头）
            for row in sheet.iter_rows(min_row=4, max_row=sheet.max_row):
                for cell in row:
                    cell.value = None

            # 列映射到Excel列
            columns_mapping = {
                "物料编号": "A",
                "层级": "C",
                "制程工段编号": "D",
                "利用率": "E",
                "用量": "F",
                "单位": "G",
                "有效起始时间": "H",
                "有效终止时间": "I",
                "上一层级物料编号": "L",
                "根节点物料编号": "M",
                "保税标识": "B"
            }

            # 写入数据
            for idx, row in data.iterrows():
                # 检查行号是否超出Excel限制
                excel_row = idx + 4  # 数据从第4行开始
                if excel_row > self.max_excel_rows:
                    self.error_occurred.emit(f"数据量超出Excel限制，将自动分割文件 (行号: {excel_row})")
                    return False
                
                for col_name, cell_column in columns_mapping.items():
                    if col_name in row:
                        sheet[f"{cell_column}{excel_row}"] = row[col_name]

            book.save(file_path)
            return True
        except Exception as e:
            error_msg = f"写入结果文件出错: {str(e)}"
            self.error_occurred.emit(error_msg)
            return False

    def split_data(self, data):
        """
        分割数据，确保相同根节点物料编号和保税标识的数据不被分开
        组内数据按层级升序排序
        """
        if data.empty:
            return []
            
        # 按根节点物料编号和保税标识分组并排序
        grouped = data.groupby(["根节点物料编号", "保税标识"], group_keys=False)
        data = grouped.apply(lambda x: x.sort_values("层级", ascending=True))
        
        # 重置索引以避免空行问题
        data = data.reset_index(drop=True)
        
        # 如果数据量小于Excel限制且不需要分割，直接返回
        if len(data) <= self.max_excel_rows - 3 and self.split_size <= 0:
            return [data]
        
        # 重新分组以确保排序后的数据保持分组
        groups = [group for _, group in data.groupby(["根节点物料编号", "保税标识"])]
        
        # 计算每个组的大小
        group_sizes = [len(group) for group in groups]
        total_size = sum(group_sizes)
        
        # 计算需要分割成多少文件
        if self.split_size > 0:
            num_files = math.ceil(total_size / min(self.split_size, self.max_excel_rows - 3))
        else:
            num_files = math.ceil(total_size / (self.max_excel_rows - 3))
            
        if num_files == 1:
            return [data]
        
        # 分割数据
        split_data = []
        current_file = pd.DataFrame()
        current_size = 0
        
        for group in groups:
            group_size = len(group)
            
            # 检查当前组是否超过单个文件限制
            if group_size > self.max_excel_rows - 3:
                error_msg = f"根节点物料 {group['根节点物料编号'].iloc[0]} (保税标识: {group['保税标识'].iloc[0]}) 的数据量({group_size})超过Excel单文件限制({self.max_excel_rows - 3})"
                self.error_occurred.emit(error_msg)
                continue
            
            # 如果当前组不能放入当前文件，则创建新文件
            if (current_size + group_size > (self.split_size if self.split_size > 0 else self.max_excel_rows - 3) 
                and current_size > 0):
                split_data.append(current_file)
                current_file = pd.DataFrame()
                current_size = 0
            
            current_file = pd.concat([current_file, group], ignore_index=True)
            current_size += group_size
        
        if not current_file.empty:
            split_data.append(current_file)
        
        return split_data

    def run(self):
        try:
            # 读取源数据
            self.message_updated.emit("正在读取源数据...")
            try:
                bom_nodes = pd.read_excel(
                    self.source_path,
                    sheet_name="LVL1以降",
                    header=2  # 第三行为表头
                )
                bom_headers = pd.read_excel(
                    self.source_path,
                    sheet_name="BOM",
                    header=0  # 第一行为表头
                )
            except Exception as e:
                error_msg = f"读取Excel文件出错: {str(e)}"
                self.error_occurred.emit(error_msg)
                raise

            self.progress_updated.emit(20)

            # 数据预处理
            self.message_updated.emit("正在预处理数据...")
            try:
                # 过滤物料编号为0的行
                bom_nodes = bom_nodes[bom_nodes["物料编号"] != 0].copy()
                bom_headers = bom_headers[bom_headers["物料编号"] != 0].copy()

                # 添加层级标记
                bom_headers["层级"] = 0

                # 用量除以1000 (仅对LVL1以降的节点)
                bom_nodes["用量"] = bom_nodes["用量"] / 1000

                # 合并数据
                combined_data = pd.concat([bom_headers, bom_nodes], ignore_index=True)

                # 筛选指定顶节点物料
                if self.top_nodes:
                    combined_data = combined_data[combined_data["根节点物料编号"].astype(str).isin(self.top_nodes)]
                    if combined_data.empty:
                        raise ValueError(f"未找到指定的顶节点物料: {', '.join(self.top_nodes)}")

                # 按照新规则进行列映射
                mapped_data = pd.DataFrame()

                # 直接映射的列
                direct_mapping = {
                    "物料编号": "物料编号",
                    "层级": "层级",
                    "制程工段编号": "制程工段编号",
                    "用量": "用量",
                    "用量单位": "单位",
                    "有效起始时间": "有效起始时间",
                    "有效终止时间": "有效终止时间",
                    "上一层级物料编号": "上一层级物料编号",
                    "根节点物料编号": "根节点物料编号"
                }

                for source_col, target_col in direct_mapping.items():
                    if source_col in combined_data.columns:
                        mapped_data[target_col] = combined_data[source_col]

                # 特殊处理的列
                # 1. 利用率固定为100.00
                mapped_data["利用率"] = 100.00
                
                # 2. Plant映射到保税标识
                if "Plant" in combined_data.columns:
                    mapped_data["保税标识"] = combined_data["Plant"].apply(
                        lambda x: "是" if x == 3402 else "否"
                    )

                # 3. 根节点单位默认为PCE
                mapped_data.loc[mapped_data["层级"] == 0, "单位"] = "PCE"

                # 处理日期
                mapped_data = self.process_dates(mapped_data)

                # 全局排序
                self.message_updated.emit("正在排序数据...")
                mapped_data = mapped_data.sort_values(
                    by=["保税标识", "根节点物料编号", "层级"],
                    ascending=[False, True, True]  # 保税标识:先是后否
                )

                # 移除可能存在的空行
                mapped_data = mapped_data.dropna(how='all')

            except Exception as e:
                error_msg = f"数据处理出错: {str(e)}"
                self.error_occurred.emit(error_msg)
                raise

            self.progress_updated.emit(60)

            # 分割数据
            self.message_updated.emit("正在分割数据...")
            data_chunks = self.split_data(mapped_data)
            total_files = len(data_chunks)
            success_files = 0

            if not data_chunks:
                raise ValueError("没有有效数据需要保存")

            # 输出到模板
            for i, chunk in enumerate(data_chunks, 1):
                self.message_updated.emit(f"正在写入文件 {i}/{total_files} (共 {len(chunk)} 行)...")
                try:
                    if total_files == 1:
                        output_filename = f"output_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.xlsx"
                    else:
                        output_filename = f"output_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}_part{i}.xlsx"
                    
                    output_path = os.path.join(self.output_dir, output_filename)
                    
                    if self.save_to_excel(chunk, output_path):
                        success_files += 1
                        self.message_updated.emit(f"成功保存文件: {output_filename}")
                except Exception as e:
                    error_msg = f"写入文件 {i} 出错: {str(e)}"
                    self.error_occurred.emit(error_msg)

            self.progress_updated.emit(100)
            self.finished_signal.emit(f"处理完成！成功保存 {success_files}/{total_files} 个文件", success_files > 0)

        except Exception as e:
            self.finished_signal.emit(f"处理过程中出错: {str(e)}", False)


class BomMergeTool(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Excel BOM 合并工具")
        self.setGeometry(300, 300, 800, 600)
        self.init_ui()

    def init_ui(self):
        # 主分割窗口，上方为操作区，下方为日志
        splitter = QSplitter(Qt.Vertical)
        
        # 上方操作区
        top_widget = QWidget()
        top_layout = QVBoxLayout()

        # 标题
        title_label = QLabel("Excel BOM 合并工具")
        title_label.setStyleSheet("font-size: 18px; font-weight: bold;")
        title_label.setAlignment(Qt.AlignCenter)
        top_layout.addWidget(title_label)

        # 文件选择表单
        form_layout = QFormLayout()
        
        # 源文件选择
        self.source_path_label = QLabel("未选择")
        self.source_path_label.setStyleSheet("color: #666;")
        source_btn = QPushButton("选择")
        source_btn.clicked.connect(self.select_source_file)
        form_layout.addRow("源文件 (ExcelA):", self.create_file_row(self.source_path_label, source_btn))

        # 模板文件选择
        self.template_path_label = QLabel("未选择")
        self.template_path_label.setStyleSheet("color: #666;")
        template_btn = QPushButton("选择")
        template_btn.clicked.connect(self.select_template_file)
        form_layout.addRow("模板文件 (ExcelB):", self.create_file_row(self.template_path_label, template_btn))

        # 输出目录选择
        self.output_path_label = QLabel("未选择")
        self.output_path_label.setStyleSheet("color: #666;")
        output_btn = QPushButton("选择")
        output_btn.clicked.connect(self.select_output_dir)
        form_layout.addRow("输出目录:", self.create_file_row(self.output_path_label, output_btn))

        # 顶节点物料筛选
        self.top_nodes_edit = QLineEdit()
        self.top_nodes_edit.setPlaceholderText("输入顶节点物料编号，多个用逗号分隔")
        form_layout.addRow("顶节点物料筛选:", self.top_nodes_edit)

        # 分割选项
        split_layout = QHBoxLayout()
        self.split_check = QComboBox()
        self.split_check.addItems(["不分割", "1000条/文件", "3000条/文件", "5000条/文件", "8000条/文件", "10000条/文件"])
        split_layout.addWidget(QLabel("分割选项:"))
        split_layout.addWidget(self.split_check)
        form_layout.addRow(split_layout)

        top_layout.addLayout(form_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        top_layout.addWidget(self.progress_bar)

        # 状态信息
        self.status_label = QLabel("准备就绪")
        self.status_label.setStyleSheet("color: #666; font-style: italic;")
        top_layout.addWidget(self.status_label)

        # 处理按钮
        process_btn = QPushButton("开始处理")
        process_btn.setStyleSheet("background-color: #4CAF50; color: white; padding: 8px;")
        process_btn.clicked.connect(self.process_files)
        top_layout.addWidget(process_btn)

        top_widget.setLayout(top_layout)
        
        # 下方日志区
        self.log_widget = QTextEdit()
        self.log_widget.setReadOnly(True)
        self.log_widget.setPlaceholderText("处理日志将显示在这里...")
        
        # 添加到分割器
        splitter.addWidget(top_widget)
        splitter.addWidget(self.log_widget)
        splitter.setSizes([300, 300])  # 设置初始大小比例
        
        self.setCentralWidget(splitter)

        # 初始化处理器
        self.processor = None

    def create_file_row(self, label, button):
        """创建文件选择行"""
        row = QWidget()
        layout = QHBoxLayout(row)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.addWidget(label)
        layout.addWidget(button)
        return row

    def select_source_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择源文件", "", "Excel 文件 (*.xlsx *.xls)"
        )
        if file_path:
            self.source_path_label.setText(file_path)

    def select_template_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择模板文件", "", "Excel 文件 (*.xlsx *.xls)"
        )
        if file_path:
            self.template_path_label.setText(file_path)

    def select_output_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择输出目录")
        if dir_path:
            self.output_path_label.setText(dir_path)

    def process_files(self):
        if not self.source_path_label.text() or self.source_path_label.text() == "未选择":
            QMessageBox.warning(self, "警告", "请选择源文件")
            return
        if not self.template_path_label.text() or self.template_path_label.text() == "未选择":
            QMessageBox.warning(self, "警告", "请选择模板文件")
            return
        if not self.output_path_label.text() or self.output_path_label.text() == "未选择":
            QMessageBox.warning(self, "警告", "请选择输出目录")
            return

        # 获取分割大小
        split_text = self.split_check.currentText()
        split_size = 0
        if split_text != "不分割":
            split_size = int(split_text.split("条")[0])

        self.progress_bar.setValue(0)
        self.status_label.setText("正在处理...")
        self.log_widget.clear()

        self.processor = BomProcessor(
            source_path=self.source_path_label.text(),
            template_path=self.template_path_label.text(),
            output_dir=self.output_path_label.text(),
            top_nodes=self.top_nodes_edit.text(),
            split_size=split_size
        )
        self.processor.progress_updated.connect(self.update_progress)
        self.processor.message_updated.connect(self.update_status)
        self.processor.finished_signal.connect(self.process_finished)
        self.processor.error_occurred.connect(self.log_error)
        self.processor.start()

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def update_status(self, message):
        self.status_label.setText(message)
        self.log_widget.append(f"[状态] {message}")

    def log_error(self, error_msg):
        """记录错误信息到日志窗口"""
        self.log_widget.append(f"[错误] {error_msg}")

    def process_finished(self, message, success):
        if success:
            self.log_widget.append(f"[完成] {message}")
            QMessageBox.information(self, "完成", message)
        else:
            self.log_widget.append(f"[错误] {message}")
            QMessageBox.critical(self, "错误", message)
        self.status_label.setText("准备就绪")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = BomMergeTool()
    window.show()
    sys.exit(app.exec_())
